#!/usr/bin/env ruby

# $Id$
#
# Copyright (c) 2007 Seggy Umboh
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# TODO: handle symbolic/hard links
# TODO: configurable/adaptive progress update frequency
# TODO: configurable/adaptive chunk size

# A tool to upload files to Amazon S3
# Requires the AWS::S3 library

require 'rubygems'
require 'aws/s3'
require 'pathname'
include AWS::S3

class UploadComplete < StandardError
end

module Net
  class HTTPGenericRequest    
    def chunk_size
      131072 # 1/8 megabyte
    end
  end
end

def track_progress(f)
  return unless $stdout.tty?
  written = 0
  pos = 0
  size = f.lstat.size
  while true
    next if f.pos == pos
    written.times { written -= $stdout.write("\b") }
    written = $stdout.write("%d%" % (pos*100/size))
    pos = f.pos
    sleep 2
  end
rescue UploadComplete
ensure
  written.times { written -= $stdout.write("\b") }
end

unless $ARGV.size == 2
  $stderr.puts <<-USAGE
    Usage: #{$0} <source> <destination>
      source: file or folder name
      destination: <bucket>/path
  
      This program reads S3 authentication information from the AMAZON_ACCESS_KEY_ID and AMAZON_SECRET_ACCESS_KEY environment variables.
  USAGE
  exit
end

unless ENV['AMAZON_ACCESS_KEY_ID'] && ENV['AMAZON_SECRET_ACCESS_KEY']
  $stderr.puts "The environment variables AMAZON_ACCESS_KEY_ID and AMAZON_SECRET_ACCESS_KEY must be set."
  exit
end

src = Pathname.new($ARGV[0])
bucket, root = $ARGV[1].split("/", 2)
root = Pathname.new(root || ".")
root += src.basename if src.directory?

$stdout.sync = true

print "Checking source..."
unless src.readable?
  $stderr.puts "Source #{src} is not readable!"
  exit
end
puts "done"

print "Connecting to Amazon S3..."
AWS::S3::Base.establish_connection!(
  :access_key_id => ENV['AMAZON_ACCESS_KEY_ID'],
  :secret_access_key => ENV['AMAZON_SECRET_ACCESS_KEY']
)
puts "done"

puts "Uploading files..."
total_files = 0
total_size = 0
src.find do |path|
  next if path.directory?
  relative_path = path.relative_path_from(src)
  target = root + relative_path
  begin
    stat = S3Object.about(target.to_s, bucket)
    if stat && stat["content-length"] == path.size.to_s && Time.parse(stat["last-modified"]) >= path.mtime
      puts "Skipping #{target}"
      next
    end
  rescue NoSuchKey
  end
  path.open do |f|
    (1..3).any? do |i|
      print "Uploading #{target}"
      print " (attempt ##{i})" if i>1
      print "..."
      t = Thread.new { track_progress(f) }
      begin
        S3Object.store(target.to_s, f, bucket)
      rescue
        t.raise(UploadComplete)
        t.join
        puts "failed: #{$!}"
        raise($!) if i==3
        sleep i
      else
        t.raise(UploadComplete)
        t.join
        puts "done"
        true
      end
    end
  end
  total_files += 1
  total_size += path.size
end
puts "Successfully uploaded #{total_files} #{total_files==1 ? 'file' : 'files'} (#{total_size/1024/1024} MB) to #{bucket}"
